package simtester.cloud

import java.security.PublicKey
import java.util.concurrent.TimeUnit
import net.schmizz.sshj.common.IOUtils
import net.schmizz.sshj.connection.channel.direct.Session
import net.schmizz.sshj.connection.channel.direct.Session.Command
import net.schmizz.sshj.SSHClient
import net.schmizz.sshj.transport.verification.HostKeyVerifier
import net.schmizz.sshj.xfer.FileSystemFile
import net.schmizz.sshj.xfer.scp.SCPFileTransfer

/**
 * Wrapper around SSHClient
 * @author linh
 *
 */
class CloudInstanceSSH {
    def address
    def user
    def authkey
    
    private test
    private verifier
    
    CloudInstanceSSH(args) {
        address = args.address
        user    = args.user
        authkey = args.authkey
        test    = args.test ?: false
        
        verifier = new HostKeyVerifier() {
            public boolean verify(String arg0, int arg1, PublicKey arg2) {
                return true; // Not bothering to verify
            }
        }
    }

    void upload(opts, file) {
        connection { ssh ->
            ssh.newSCPFileTransfer().upload(new FileSystemFile(file.absolutePath), opts.to)
            //TODO: Catch an IOException?
        }
    }
    
    def exec(commandString) {
        boolean allocatePty = false
        if (commandString =~ /^sudo /) {
            allocatePty = true
        }
        
        def output
        session(allocatePty) { s ->
            output = execWithOutput(session: s, cmd: commandString)
        }
        
        return output
    }
    
    void execAsync(commandString) {
        session(false) { s ->
            execWithNoOutput(session: s, cmd: commandString)
        }
    }

    
    private def execWithOutput(args) {
        def timeout = args.timeout ?: 120

        final Command cmd = args.session.exec(args.cmd)
        cmd.join(timeout, TimeUnit.SECONDS)
        def output = IOUtils.readFully(cmd.getInputStream()).toString().trim()
        
        return output
    }
    
    private void execWithNoOutput(args) {
        args.session.exec(args.cmd)
    }

    private def tryToConnect(ssh, address) {
        def timeout  = test ? 60 : 120000
        def timeincr = test ?  5 :   5000
        def lastError
        
        try {
            Utils.waitFor(limit: timeout, step: timeincr) {
                try {
                    ssh.connect(address)
                    return true  // Break out of wait if connect is successful
                } catch (ConnectException e) {
                    lastError = e
                } catch (NoRouteToHostException e) {
                    lastError = e
                }
            }
        } catch (TimeoutException e) {
            throw new CloudException("Cannot connect to $address: ${lastError.message}")
        }
    }
    
    private def connection(Closure closure) {
        final SSHClient ssh = new SSHClient()
        ssh.addHostKeyVerifier(verifier)
        
        tryToConnect(ssh, address)
        //TODO: Handle case where connect is not successful
        
        ssh.authPublickey(user, authkey)
        
        try {
            closure.call(ssh)
        } finally {
            ssh.disconnect()
        }
    }
    
    private void session(boolean allocatePty = false, Closure closure) {
        connection { ssh ->
            def session = ssh.startSession()
            
            if (allocatePty) {
                session.allocateDefaultPTY()
            }
            
            try {
                closure.call(session)
            } finally {
                session.close()
            }
        }
    }
}